#! /usr/bin/env python3

# 1.导包
import rospy
from superglue.srv import SuperScore, SuperScoreRequest, SuperScoreResponse
from cv_bridge import CvBridge
import cv_bridge
print(' - cv_bridge.__file__ = ',cv_bridge.__file__)
from sensor_msgs.msg import Image
import sensor_msgs


import cv2
import numpy as np
import torch
import yaml


import sys
sys.path.append("/home/daybeha/Documents/github/DeepLabV3_ws/src/superglue")
from Detectors import create_detector
from Matchers import create_matcher

from models.utils import AverageTimer
from utils.tools import *


class Matching(torch.nn.Module):
    """ Image Matching Frontend (SuperPoint + SuperGlue) """
    def __init__(self, config={}):
        super().__init__()
        # create detector
        self.detector = create_detector(config["detector"])
        # create matcher
        self.matcher = create_matcher(config["matcher"])


        # self.superpoint = SuperPoint(config.get('superpoint', {}))
        # self.superglue = SuperGlue(config.get('superglue', {}))

    def forward(self, data):
        """ Run SuperPoint (optionally) and SuperGlue
        SuperPoint is skipped if ['keypoints0', 'keypoints1'] exist in input
        Args:
          data: dictionary with minimal keys: ['image0', 'image1']
        """
        pred = {'ref': None, 'cur': None}

        # TODO 这块的显存存占用有待优化
        # Extract SuperPoint (keypoints, scores, descriptors) if not provided
        if 'keypoints0' not in data:
            pred['ref'] = self.detector(data['image0'])
        if 'keypoints1' not in data:
            pred['cur'] = self.detector(data['image1'])

        matches = self.matcher(pred)

        return matches

# 这一句至关重要！！！ 能节省至少一半显存！！！
torch.set_grad_enabled(False)

config = "/home/daybeha/Documents/github/DeepLabV3_ws/src/superglue/params/superpoint_supergluematch.yaml"
with open(config, 'r') as f:
    config = yaml.safe_load(f)


# print(f"cuda avaliable: {torch.cuda.is_available()}")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("device: ", device)
config["device"] = device
matching = Matching(config).eval().to(device)


def compute_score(img0, img1):
    matches = matching({'image0': img0, 'image1': img1})
    score = np.mean(matches["match_score"].cpu().detach().numpy())
    return score


def con_show(img0, img1):
    # 纵向连接 image = np.vstack((img0, img1))
    # 横向连接 image = np.concatenate([img0, img1], axis=1)
    image = np.concatenate((img0, img1))

    cv2.imshow("image in python", image)
    # cv2.imwrite("/home/tt/test/111.png", img)
    cv2.waitKey(0)


def show_img(img=None):
    if img is not None:
        cv2.imshow("image in python", img)
        # cv2.imwrite("/home/tt/test/111.png", img)
        cv2.waitKey(1)



# 回调函数的参数是请求对象，返回值是响应对象
def imgCallback(req):
    rospy.loginfo(f"received request")


    bridge = CvBridge()
    # img0 = bridge.imgmsg_to_cv2(req.image0, "bgr8")
    # img1 = bridge.imgmsg_to_cv2(req.image1, "bgr8")
    img0 = bridge.imgmsg_to_cv2(req.image0, "mono8")
    img1 = bridge.imgmsg_to_cv2(req.image1, "mono8")

    try:
        # con_show(img0, img1)
        # show_img(img0)
        score = compute_score(img0, img1)
    except:
        rospy.sleep(0.2)
        print("WARNING: out of memory")
        if hasattr(torch.cuda, 'empty_cache'):
            torch.cuda.empty_cache()

        score = compute_score(img0, img1)

    return SuperScoreResponse(score)


if __name__ == "__main__":

    # 2.初始化 ROS 节点
    rospy.init_node("superglue_server")
    # 3.创建服务对象
    server = rospy.Service("/super_score", SuperScore, imgCallback)
    # 4.回调函数处理请求并产生响应


    # 5.spin 函数
    rospy.spin()





